#!/usr/bin/env python
#-*- coding: UTF-8 -*-
# nodmgr.py: node manager core implementaion
#
# Copyright (C) 2008-2010 Hily Jiang <hilyjiang@gmail.com>
#
# This file is part of dtask.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

"""
Node Manager core implementaion.
"""

import sys
import threading
import sqlite3
from time import time


class NodeManager(object):
    """Node Manager, it hold a node list."""

    def __init__(self, dbpath=':memory:', jobsPerNode=2, encoding='utf-8'):
        """
        dbpath:
            sqlite3 database save path, default ':memory:' means run in memory.
        jobsPerNode:
            max jobs per node can process at the same time
        encoding:
            default encoding of sqlite3 output is unicode, we always use utf-8.
        """
        self.dbpath = dbpath
        self.jobsPerNode = jobsPerNode
        self.encoding = encoding

        # blocking implemetation resource
        self.idleJobs = 0
        self.idleJobEvent = threading.Event()

        # database operation implementation resource
        self.dbThreadLock = threading.Semaphore(0)
        self.dbThreadLock.release()
        self.dbOpEvent = threading.Event()
        self.dbReturnEvent = threading.Event()
        self.dbOpFunc = None
        self.dbOpArgs = None
        self.dbOpKwrgs = None
        self.dbReturn = None

        # Because sqlite3 does not support cross-thread connection sharing,
        # so we create an individual thread to handle sqlite3 operations.
        self.dbthread = threading.Thread(target=self.__dbThreadFunc)
        self.dbthread.setDaemon(True)
        self.dbthread.start()

    def __dbThreadFunc(self):
        """Thread function for sqlite3 operations."""
        self.conn = sqlite3.connect(self.dbpath)
        self.cur = self.conn.cursor()
        try:
            # create table for node manager
            # nodid:     node ID
            # nodname:   node name
            # nodtype:   node type
            # nodip:     node ip
            # nodport:   node port
            # status:    node status
            # curjobs:   count of jobs currently hold
            # totaljobs: count of jobs totally processed
            # logintime: login time
            # lastfinishtime: last time finish a task
            self.cur.execute('''
                CREATE TABLE nodes(
                    nodid INTEGER PRIMARY KEY AUTOINCREMENT,
                    nodname CHAR(20),
                    nodtype CHAR(20),
                    nodip CHAR(20),
                    nodport INTEGER,
                    status CHAR(20),
                    curjobs INTEGER,
                    totaljobs INTEGER,
                    logintime INTEGER,
                    lastfinishtime INTEGER
                )
            ''')
            self.conn.commit()
        except Exception, e: # table exists
            self.cur.execute('DELETE FROM nodes')
            self.conn.commit()
            print '[MGR]', e

        while True:
            self.dbOpEvent.wait()
            self.dbOpEvent.clear()
            self.dbReturn = self.dbOpFunc(self, *self.dbOpArgs,
                                          **self.dbOpKwargs)
            self.dbReturnEvent.set()

    def __dbThread(func):
        """Decoration function to wrap functions that operate sqlite3.
           The functions decorated will then executed in the db thread."""
        def wrapFunc(self, *args, **kwargs):
            self.dbThreadLock.acquire()
            self.dbOpFunc = func
            self.dbOpArgs = args
            self.dbOpKwargs = kwargs
            self.dbOpEvent.set()
            self.dbReturnEvent.wait()
            rt = self.dbReturn
            self.dbReturnEvent.clear()
            self.dbThreadLock.release()
            return rt
        return wrapFunc

    @__dbThread
    def newNode(self, nodIp, nodPort, nodType, nodName):
        """Add a new node to node manager. Return an ID identifies the node."""
        nodPort = int(nodPort)
        nowTime = int(time())
        sql = 'INSERT INTO nodes VALUES(NULL, ?, ?, ?, ?, "normal", 0, 0, ?, 0)'
        self.cur.execute(sql, (nodName, nodType, nodIp, nodPort, nowTime))
        self.conn.commit()
        rowid = self.cur.lastrowid
        sql = 'SELECT nodid from nodes WHERE ROWID=?'
        nodeId = self.cur.execute(sql, (rowid,)).fetchone()[0]

        self.idleJobs += self.jobsPerNode
        self.idleJobEvent.set()

        # return row id as node id
        return nodeId

    @__dbThread
    def __acquireNodeFromDB(self):
        """Get a available node from database."""
        self.cur.execute('SELECT nodid FROM nodes ' \
                         ' WHERE curjobs<? AND status="normal"' \
                         ' ORDER BY curjobs, lastfinishtime',
                         (self.jobsPerNode,))
        nodeId = self.cur.fetchone()[0]
        self.cur.execute('UPDATE nodes SET curjobs=curjobs+1,' \
                         ' totaljobs=totaljobs+1 WHERE nodid=?', (nodeId,))
        self.conn.commit()
        return nodeId

    def acquireNode(self):
        """Get a available node from list. Return node ID."""
        while self.idleJobs <= 0:
            self.idleJobEvent.wait()
            self.idleJobEvent.clear()
        self.idleJobs -= 1
        if self.idleJobs > 0:
            self.idleJobEvent.set()
        return self.__acquireNodeFromDB()

    @__dbThread
    def releaseNode(self, nodeId):
        """Release the available node to list."""
        self.cur.execute('UPDATE nodes SET curjobs=curjobs-1,' \
                         ' lastfinishtime=? WHERE nodid=?',
                         (int(time()), nodeId))
        self.conn.commit()
        self.idleJobs += 1
        self.idleJobEvent.set()

    @__dbThread
    def getNodeInfo(self, nodeId):
        """Get node info."""
        nodeId = int(nodeId)
        self.cur.execute('SELECT nodname, nodip, nodport FROM nodes' \
                         ' WHERE nodid=?', (nodeId,))
        return self.cur.fetchone()

    @__dbThread
    def removeNode(self, nodeId):
        """Remove a node from list."""
        self.cur.execute('SELECT curjobs FROM nodes WHERE nodid=?',(nodeId,))
        curjobs = self.cur.fetchone()[0]
        self.idleJobs -= self.jobsPerNode - curjobs
        self.cur.execute('DELETE FROM nodes WHERE nodid=?', (nodeId,))
        self.conn.commit()

